[Spyre-Next] [Cleanup] Rework of CustomOp wrapping#842
[Spyre-Next] [Cleanup] Rework of CustomOp wrapping#842bohnstingl merged 12 commits intotorch-spyre:mainfrom
Conversation
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
|
👋 Hi! Thank you for contributing to vLLM support on Spyre. We also recommend installing prek and configuring it to check your code before every local commit. |
github-actions
Bot
changed the title
|
bot:next-test |
| return prefix | ||
|
|
||
|
|
||
| def create_rmsnorm_op_pair(): |
There was a problem hiding this comment.
what's the rationale for moving op-specific methods to a generic utilities package? I would expect this to continue living in rms_norm.py as e.g. rms_norm.create_op_pair()
There was a problem hiding this comment.
Yeah, I was debating about this myself. The issue here is that although the two functions are really close, upstream vLLM runs the infer_schema function, which requires explicit knowledge about the inputs and those inputs are specific for the CustomOp at hand. I left the _fake_impl function in utils.py, but the op specific function I moved back to the individual files.
|
^^ looks like tests are failing, mind taking a look into it? |
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>
|
As a general comment, the CustomOp infrastructure will be overhauled once vLLM has its vLLM IR landed. I looked at it and I think we should be in a good position for a quick turn-over once the vLLM IR lands. In fact, the refactor of this PR should simplify this change even further |
Signed-off-by: Joe Runde <joe@joerun.de>
| return pytree.tree_map( | ||
| lambda el: el[:orig_batch_size, :], | ||
| convert_from_spyre(outs, dtype=x_dtype, device=x_device), | ||
| lambda el: convert(el, dtype=x_dtype, device=x_device)[:orig_batch_size, :], |
There was a problem hiding this comment.
I think there was a small bug here where convert wasn't in the transform so in the case where there is a residual, a tuple was being passed to convert which started failing with the new unit tests
|
bot:next-test |
joerunde
left a comment
joerunde
left a comment
There was a problem hiding this comment.
Nice cleanup!
I've merged in main and 🤞 the tests should be passing and we'll be good to merge
|
Tests are passing! @bohnstingl feel free to merge if you're fine with the tiny bugfix I added on the output conversion in rms_norm.py |
2 participants
Description
This PR intends to simplify the way CustomOps are currently wrapped for execution on
spyre. In addition to a simplification, this rework enables more features from upstream vLLM, for example the compilation interface. This will eventually be needed when larger parts of the model is executed on spyre.Related Issues
#733
Test Plan
This change should be transparent to users and thus the already existing/running tests should not change.
Checklist
bash format.sh)Signed-off-by:line (DCO compliance)